package net.cyue.ort.llm.data;

import ai.onnxruntime.OnnxTensor;
import ai.onnxruntime.OnnxValue;

import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Beam搜索候选类
 */
public class BeamCandidate implements Comparable<BeamCandidate>, AutoCloseable {
    private final TokenList tokenList;
    private double sequenceScore;
    private Map<String, OnnxTensor> cache;
    private double diversity = 1.0;
    private double normalizedScore;
    
    public BeamCandidate(
        Token token,
        double sequenceScore,
        Map<String, OnnxTensor> cache
    ) {
        this.tokenList = new TokenList();
        this.sequenceScore = sequenceScore;
        this.cache = cache;
        if (token != null) {
            this.tokenList.add(token);
        }
    }
    
    public BeamCandidate(
        List<Token> tokens,
        double sequenceScore,
        Map<String, OnnxTensor> cache
    ) {
        this.tokenList = new TokenList(tokens);
        this.sequenceScore = sequenceScore;
        this.cache = cache;
    }
    
    public BeamCandidate(
        TokenList tokenList,
        double sequenceScore,
        Map<String, OnnxTensor> cache
    ) {
        this.tokenList = tokenList != null ? tokenList : new TokenList();
        this.sequenceScore = sequenceScore;
        this.cache = cache;
    }
    
    public Map<String, OnnxTensor> getCache() {
        return this.cache;
    }
    
    public double getSequenceScore() {
        return this.sequenceScore;
    }
    
    public double getNormalizedScore() {
        return this.normalizedScore;
    }
    
    public void setNormalizedScore(double normalizedScore) {
        this.normalizedScore = normalizedScore;
    }
    
    @Override
    public int compareTo(BeamCandidate other) {
        return Double.compare(other.getNormalizedScore(), this.getNormalizedScore());
    }
    
    @Override
    public void close() {
        cleanup();
    }
    
    @Override
    public String toString() {
        return "BeamCandidate {" +
            "\n\ttokens=" + tokenList +
            ",\n\tsequenceScore=" + sequenceScore +
            ",\n\tdiversity=" + diversity +
            "\n}";
    }
    
    public void cleanup() {
        if (cache != null) {
            cache.values().stream()
                .filter(tensor -> tensor != null && !tensor.isClosed())
                .forEach(OnnxValue::close);
            cache.clear();
        }
    }
    
    public int getTokenCount() {
        return tokenList.size();
    }
    
    public BeamCandidate plus(BeamCandidate candidate) {
        this.tokenList.addAll(candidate.tokenList);
        double decayFactor = Math.pow(0.95, this.tokenList.size());
        this.sequenceScore += candidate.sequenceScore * decayFactor;
        if (this.cache != null) {
            this.cache.values().stream()
                .filter(tensor -> tensor != null && !tensor.isClosed())
                .forEach(OnnxValue::close);
        }
        this.cache = candidate.cache;
        return this;
    }
    
    public List<Long> getTokenIDList() {
        return tokenList.getTokenIDList();
    }
    
    public void updateDiversity(List<BeamCandidate> otherCandidates) {
        if (otherCandidates == null || otherCandidates.isEmpty()) {
            return;
        }
        
        Set<Long> thisTokenSet = new HashSet<>(this.getTokenIDList());
        double totalOverlap = otherCandidates.stream()
            .filter(other -> other != this)
            .mapToDouble(other -> {
                Set<Long> otherTokenSet = new HashSet<>(other.getTokenIDList());
                Set<Long> intersection = new HashSet<>(thisTokenSet);
                intersection.retainAll(otherTokenSet);
                Set<Long> union = new HashSet<>(thisTokenSet);
                union.addAll(otherTokenSet);
                return union.isEmpty() ? 0 : (double) intersection.size() / union.size();
            })
            .sum();
        
        this.diversity = 1.0 - (totalOverlap / otherCandidates.size());
    }
    
    public TokenList getTokenList() {
        return tokenList;
    }
}

